package club.babyf.nativeai.service.impl;

import club.babyf.nativeai.config.NativeAiProperties;
import club.babyf.nativeai.service.IEmbeddingService;
import club.babyf.nativeai.utils.JsonCustomUtil;
import com.fasterxml.jackson.core.type.TypeReference;
import com.fasterxml.jackson.databind.JsonNode;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang.ArrayUtils;
import org.springframework.stereotype.Component;

import java.net.URI;
import java.net.http.HttpClient;
import java.net.http.HttpRequest;
import java.net.http.HttpResponse;
import java.time.Duration;
import java.util.List;
import java.util.Map;

/**
 * EmbeddingServiceImpl.
 *
 * @author zhanglei.
 * @date 2025/3/17 15:46.
 * @description 向量服务实现.
 */
@Slf4j
@Component
@RequiredArgsConstructor
public class EmbeddingServiceImpl implements IEmbeddingService {

    private final NativeAiProperties nativeAiProperties;

    private final HttpClient httpClient = HttpClient.newBuilder().connectTimeout(Duration.ofSeconds(300)).build();

    @Override
    public float[] embedding(String text) {
        try {
            String requestBody = JsonCustomUtil.toJsonString(Map.of(
                    "model", nativeAiProperties.getOpenai().getEmbeddingModel(),
                    "input", text
            ));

            HttpRequest request = buildRequest(nativeAiProperties.getOpenai().getBaseUrl() + "/v1/embeddings", requestBody);
            HttpResponse<String> response = httpClient.send(request, HttpResponse.BodyHandlers.ofString());

            return parseEmbeddingResponse(response.body());
        } catch (Exception e) {
            log.error("向量生成失败", e);
            throw new RuntimeException("OpenAI服务调用失败");
        }
    }

    /**
     * 构建通用请求
     *
     * @param url  请求URL
     * @param body 请求体
     * @return 构建的HttpRequest对象
     */
    private HttpRequest buildRequest(String url, String body) {
        return HttpRequest.newBuilder()
                .uri(URI.create(url))
                .header("Content-Type", "application/json")
                .header("Authorization", "Bearer " + nativeAiProperties.getOpenai().getApiKey())
                .POST(HttpRequest.BodyPublishers.ofString(body))
                .timeout(Duration.ofSeconds(nativeAiProperties.getOpenai().getTimeoutSec()))
                .build();
    }

    /**
     * 解析向量响应
     *
     * @param json 响应JSON字符串
     * @return 解析后的向量数组
     */
    private float[] parseEmbeddingResponse(String json) {
        JsonNode root = JsonCustomUtil.toJsonNode(json);
        List<Float> floatList = JsonCustomUtil.convert(root.path("data").get(0).path("embedding"), new TypeReference<List<Double>>() {
                })
                .stream()
                .map(Double::floatValue)
                .toList();
        return ArrayUtils.toPrimitive(floatList.toArray(new Float[0]));
    }

}
